import mb4_file
import SHA2_RSA2048_file



#这是命令行版本，写都写了，就不删了
def mb4_init():
    while(1):
        print("请选择你的加密模式:\n1.sm4加密\n2.sm4解密\n0.退出")
        mb4_model=input()
        match mb4_model:
            case "1":
                print("请输入你的密钥:")
                str=input()
                key=mb4_file.get_key(str)
                # print("请输入你的iv初始向量:")
                iv=mb4_file.get_iv(str)
                print("请输入你需要加密的文件路径:")
                plaintext_bytes=mb4_file.exp_file(input())
                key_bytes=bytes(key,encoding="utf8")
                iv_bytes=bytes(iv,encoding="utf8")
                ciphertext_bytes=mb4_file.SM4_cbc_enc(plaintext_bytes,iv_bytes,key_bytes)
                print("请选择要保存的文件路径以及文件名:")
                url=input()
                with open(url,'wb')as f:
                    #f.truncate
                    f.write(ciphertext_bytes)
                f.close
            case "2":
                print("请输入你的密钥：")
                str=input()
                key=mb4_file.get_key(str)
                # print("请输入你的iv初始向量:")
                iv=mb4_file.get_iv(str)
                # if(len(key)+len(iv)<32):
                #     print("密钥和初始向量不足16位")
                #     return
                key_bytes=bytes(key,encoding="utf8")
                iv_bytes=bytes(iv,encoding="utf8")
                print("请输入你需要解密的文件路径:")
                ciphertext_bytes=mb4_file.exp_file(input())
                plaintext_bytes=mb4_file.SM4_cbc_dec(ciphertext_bytes,iv_bytes,key_bytes)
                print("请选择要保存的文件路径以及文件名:")
                url=input()
                with open(url,'wb')as f:
                    #f.truncate
                    f.write(plaintext_bytes)
                f.close
            case _:
                return


def RSA_init():
    while(1):
        print("请选择你的加密模式:\n1.生成RSA公私钥\n2.为文件创建数字证书并签字\n3.验证文件中数字证书")
        RSA_model=input()
        match RSA_model:
            case "1":
                print("请选择要保存的公钥文件路径以及文件名:")
                pub_url=input()
                print("请选择要保存的私钥文件路径以及文件名:")
                pri_url=input()
                SHA2_RSA2048_file.gen_rsa_key(pri_url,pub_url)
            case "2":
                print("请输入的私钥文件路径以及文件名")
                private_key=SHA2_RSA2048_file.exp_file(input())
                # with open(input(),"rb")as f:
                #     private_key=f.read()
                # f.close()
                print("请输入你需要加密的文件路径:")
                plaintext=SHA2_RSA2048_file.exp_file(input())
                signature=SHA2_RSA2048_file.rsa_sign(plaintext,private_key)
                print("请选择要保存的证书文件路径以及文件名:")
                with open(input(),"wb") as f:
                    f.write(signature)
                f.close()
            case "3":
                print("请输入的公钥文件路径以及文件名")
                public_key=SHA2_RSA2048_file.exp_file(input())
                print("请输入你需要验证的文件路径:")
                plaintext=SHA2_RSA2048_file.exp_file(input())
                print("请输入你需要验证的证书路径:")
                signature=SHA2_RSA2048_file.exp_file(input())
                result=SHA2_RSA2048_file.rsa_verify(signature,plaintext,public_key)
                print("验证结果为:",result)
            case _:
                return
                



while(1):
    print("--加密系统--")
    print("请选择你的加密模式:\n1.sm4加密解密\n2.加密验证证书\n0.退出系统")
    model=input()
    match model:
        case "1":
            mb4_init()
        case "2":
            RSA_init()
        case "0":
            exit(0)
        case _:
            print("你在输入错误的模式，正在退出系统")
            exit(0)
